# AIWeb_fast/mcpTool/useMCPAgent.py
import asyncio
import datetime
import json
import logging
import traceback
from typing import Dict, Any, List

from dotenv import load_dotenv

from config.FileHandle import load_mcp_tools_config
from llm import llmConfig
from llm.prompt import getPrompt
from mcpTool.MCPClientManager import MCPClientManager
from mcpTool.mcpCheck import prepare_environment
from mcpTool.toolUtils import format_tool_response, select_tool_with_llm_sync

# 加载环境变量
load_dotenv()

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger('MCP-Agent')

# 全局MCP管理器引用 - 将从app.main中获取实例
manager = None


async def initialize_mcp_manager():
    """初始化全局MCP管理器引用"""
    global manager
    if manager is None:
        manager = MCPClientManager()
        await manager.initialize()
        logger.info("成功获取到app.main中的MCP管理器实例")
    return manager


async def mcp_usage_async(message: str, tool_config: Dict[str, Any], session_id: str = "default", max_retries: int = 1):
    """使用MCP工具处理用户消息(异步版本)

    Args:
        message: 用户输入的消息
        tool_config: MCP工具配置
        session_id: 会话ID
        max_retries: 最大重试次数

    Yields:
        工具调用过程中的日志和最终结果
    """
    # 准备环境
    # prepare_environment(tool_config)

    # 确保MCP管理器已初始化
    global manager
    await initialize_mcp_manager()

    retry_count = 0
    last_error = None

    while retry_count <= max_retries:
        try:
            # 根据连接类型选择连接方式
            connection_type = tool_config.get("command", "stdio")

            if connection_type == "sse":
                # 使用SSE连接
                base_url = tool_config.get("base_url")
                if not base_url:
                    raise Exception("使用SSE连接时必须提供base_url")

                sse_session_id = session_id

                yield json.dumps({
                    "type": "message",
                    "content": f"使用SSE连接到MCP服务器: {base_url} (session_id: {sse_session_id})\n",
                    "timestamp": datetime.datetime.now().isoformat()
                })

                # 尝试连接，最多重试3次
                connect_retries = 3
                connect_count = 0
                while connect_count < connect_retries:
                    success = await manager.connect_sse(base_url, sse_session_id)
                    if success:
                        break
                    connect_count += 1
                    if connect_count < connect_retries:
                        yield json.dumps({
                            "type": "message",
                            "content": f"SSE连接失败，正在重试 ({connect_count}/{connect_retries})...\n",
                            "timestamp": datetime.datetime.now().isoformat()
                        })
                        await asyncio.sleep(1)  # 等待1秒后重试
                    else:
                        raise Exception(f"无法使用SSE连接到MCP服务器: {base_url}")

            else:
                # 使用命令行方式连接
                command = tool_config.get("command")
                args = tool_config.get("args", [])
                env = tool_config.get("env", {})

                # 检查是否需要切换连接
                if not manager.is_connected or (manager.current_connection and
                                                (manager.current_connection.get("command") != command or
                                                 manager.current_connection.get("args") != args)):

                    yield json.dumps({
                        "type": "message",
                        "content": f"连接到MCP服务器: {command} {' '.join(args or [])}\n",
                        "timestamp": datetime.datetime.now().isoformat()
                    })
                    success = await manager.connect(command, args, env)
                    if not success:
                        raise Exception(f"无法连接到MCP服务器: {command} {' '.join(args or [])}")

            # 列出可用工具
            tools = await manager.list_tools()
            yield json.dumps({
                "type": "message",
                "content": f"可用工具数量: {len(tools)}\n",
                "timestamp": datetime.datetime.now().isoformat()
            })
            
            # 输出工具具体内容，使用序号格式
            tool_list_content = "可用工具列表:\n"
            for i, tool in enumerate(tools, 1):
                # 尝试获取更多工具信息
                tool_name = tool.name if hasattr(tool, 'name') else f"工具{i}"
                tool_description = tool.description if hasattr(tool, 'description') else "无描述"
                
                # 基本信息
                tool_info = f"{i}. {tool_name}: {tool_description}"
                
                # 添加参数信息（如果有）
                try:
                    if hasattr(tool, 'parameters') and tool.parameters:
                        params = tool.parameters
                        # 检查参数是否为字典类型
                        if isinstance(params, dict) and params:
                            tool_info += "\n   参数:"
                            for param_name, param_info in params.items():
                                # 尝试安全地提取参数描述
                                if isinstance(param_info, dict) and 'description' in param_info:
                                    param_desc = param_info['description']
                                elif hasattr(param_info, 'description'):
                                    param_desc = getattr(param_info, 'description', '')
                                else:
                                    param_desc = str(param_info)
                                
                                # 提取是否必需
                                required = ""
                                if isinstance(param_info, dict) and 'required' in param_info:
                                    required = " (必需)" if param_info['required'] else " (可选)"
                                elif hasattr(param_info, 'required'):
                                    required = " (必需)" if getattr(param_info, 'required') else " (可选)"
                                
                                tool_info += f"\n   - {param_name}{required}: {param_desc}"
                except Exception as e:
                    logger.warning(f"处理工具参数时出错: {str(e)}")
                
                tool_list_content += f"{tool_info}\n\n"
            
            # 输出序号格式的工具列表
            yield json.dumps({
                "type": "message",
                "content": tool_list_content,
                "timestamp": datetime.datetime.now().isoformat()
            }, ensure_ascii=False)
            
            # 使用LLM选择合适的工具和参数
            name, arguments = select_tool_with_llm_sync(tools, message)
            yield json.dumps({
                "type": "message",
                "content": f"选择调用工具: {name} 参数: {arguments}\n",
                "timestamp": datetime.datetime.now().isoformat()
            })

            # 调用工具
            response = await manager.call_tool(name, arguments)

            # 处理返回结果，确保返回格式统一
            result = format_tool_response(response, False, "工具调用成功")

            # 断开连接
            if manager.is_connected:
                await manager.disconnect()

            # 尝试从返回结果中提取有效内容
            extracted_content = "获取的内容:"
            try:
                # 检测是否为TextContent对象
                if isinstance(result, dict) and 'content' in result:
                    content_data = result['content']
                    
                    # 处理内容为TextContent类型的情况
                    if hasattr(content_data, '__dict__') and hasattr(content_data, 'text'):
                        json_text = content_data.text
                        extracted_content += f"\n{json_text}"
                    elif isinstance(content_data, list):
                        # 处理content为列表的情况
                        for i, item in enumerate(content_data, 1):
                            if hasattr(item, 'text'):
                                # TextContent对象
                                extracted_content += f"\n{i}. {item.text}"
                            elif isinstance(item, dict) and 'text' in item:
                                # 包含text字段的字典
                                extracted_content += f"\n{i}. {item['text']}"
                            else:
                                # 其他类型
                                extracted_content += f"\n{i}. {str(item)}"
                    else:
                        # 其他格式，尝试JSON编码
                        extracted_content += f"\n```json\n{json.dumps(content_data, ensure_ascii=False, indent=2)}\n```"
                else:
                    # 普通对象，直接JSON编码
                    extracted_content += f"\n```json\n{json.dumps(result, ensure_ascii=False, indent=2)}\n```"
            except Exception as e:
                logger.warning(f"提取内容时出错: {str(e)}")
                extracted_content += f"\n{result}"
            
            # 返回处理后的结果
            yield json.dumps({
                "type": "message",
                "content": f"{extracted_content}\n",
                "timestamp": datetime.datetime.now().isoformat()
            })

            # 检查错误状态，如果有错误则提前结束
            is_error = False
            error_message = ""

            # 检查外层isError
            if isinstance(result, dict) and result.get("isError", False):
                is_error = True
                error_message = result.get("message", "工具调用失败")
            # 检查内层content中的isError
            elif (isinstance(result, dict) and
                  isinstance(result.get("content"), dict) and
                  result.get("content", {}).get("isError", False)):
                is_error = True
                # 尝试从content中获取错误信息
                content_data = result.get("content", {})
                if isinstance(content_data.get("content"), list) and len(content_data.get("content", [])) > 0:
                    # 尝试获取第一个内容项的text
                    first_item = content_data["content"][0]
                    if hasattr(first_item, "text"):
                        error_message = first_item.text
                    elif isinstance(first_item, dict) and "text" in first_item:
                        error_message = first_item["text"]
                    else:
                        error_message = str(first_item)
                else:
                    error_message = content_data.get("message", "内容处理失败")

            # 如果存在错误，直接返回错误消息并结束
            if is_error:
                yield json.dumps({
                    "type": "result",
                    "content": f"**处理出错**\n\n{error_message}",
                    "timestamp": datetime.datetime.now().isoformat()
                })
                return

            # 使用大模型处理结果并转换为 Markdown 格式
            llm = llmConfig.getCommonLLM(streaming=True)
            
            # 构建提示词
            prompt = getPrompt("markDownHandle", message, result)
            print(f"Prompt: {prompt}")
            
            # 初始化一个空字符串来存储完整响应
            full_response = ""
            
            try:
                # 使用流式模式调用大模型
                async for chunk in llm.astream(prompt):
                    try:
                        if chunk:
                            # 获取chunk的内容
                            chunk_content = chunk.content if hasattr(chunk, 'content') else str(chunk)
                            # 累积完整响应
                            full_response += chunk_content
                            
                            # 实时返回每个片段给前端
                            yield json.dumps({
                                "type": "result",
                                "content": chunk_content,
                                "timestamp": datetime.datetime.now().isoformat(),
                            })
                    except Exception as chunk_error:
                        logger.error(f"处理流式响应片段时出错: {str(chunk_error)}")
                        # 继续处理下一个片段
                        continue
                
                # 流式输出结束，发送完成标记
                yield json.dumps({
                    "type": "result",
                    "content": "",
                    "timestamp": datetime.datetime.now().isoformat(),
                })
                
                print(f"Processed Result: {full_response}")
            except Exception as e:
                logger.error(f"流式处理过程中出错: {str(e)}")
                # 发送错误消息
                yield json.dumps({
                    "type": "result",
                    "content": f"**处理出错**\n\n流式处理过程中发生错误: {str(e)}",
                    "timestamp": datetime.datetime.now().isoformat(),
                    "error": True
                })
            finally:
                return

        except Exception as e:
            last_error = e
            error_detail = traceback.format_exc()
            yield json.dumps({
                "type": "result",
                "content": f"MCP工具调用失败 (尝试 {retry_count + 1}/{max_retries + 1}): {str(e)}\n",
                "timestamp": datetime.datetime.now().isoformat()
            })

            retry_count += 1
            if retry_count <= max_retries:
                yield json.dumps({
                    "type": "message",
                    "content": "尝试重新连接MCP服务器...\n",
                    "timestamp": datetime.datetime.now().isoformat()
                })
                await asyncio.sleep(1)  # 等待1秒后重试
                continue
        finally:
            # 断开连接
            if manager.is_connected:
                await manager.disconnect()

    # 达到最大重试次数后仍然失败
    error_result = format_tool_response(None, True, f"工具调用失败: {str(last_error) or '未知错误，请检查服务器日志'}")
    yield json.dumps({
        "type": "result",
        "content": f"{error_result}\n",
        "timestamp": datetime.datetime.now().isoformat()
    })


async def use_MCP_tools_async(message: str, tool_ids: List[str], session_id: str = "default"):
    """
    使用MCP工具处理用户消息(异步版本)

    Args:
        message: 用户输入的消息
        tool_ids: 要使用的MCP工具ID列表
        session_id: 会话ID

    Yields:
        包含处理结果的字典
    """
    try:
        if not tool_ids:
            yield json.dumps({
                "type": "message",
                "content": "未提供工具ID，请指定要使用的MCP工具ID\n",
                "timestamp": datetime.datetime.now().isoformat()
            })
            return

        # 读取工具配置
        tools_config = load_mcp_tools_config(tool_ids)
        if not tools_config:
            yield json.dumps({
                "type": "message",
                "content": "未找到有效的MCP工具配置\n",
                "timestamp": datetime.datetime.now().isoformat()
            })
            return

        # 获取第一个工具配置
        tool_config = tools_config[0]
        yield json.dumps({
            "type": "message",
            "content": f"工具配置: command={tool_config.get('command')}, args={tool_config.get('args')}\n",
            "timestamp": datetime.datetime.now().isoformat()
        })

        # 调用MCP工具
        yield json.dumps({
            "type": "message",
            "content": f"调用工具: {tool_config.get('name', '未知工具')}\n",
            "timestamp": datetime.datetime.now().isoformat()
        })

        async for chunk in mcp_usage_async(message, tool_config, session_id):
            yield chunk

    except Exception as e:
        error_detail = traceback.format_exc()
        logger.error(f"MCP工具执行失败: {str(e) or '未知错误'}")
        logger.debug(error_detail)
        error_message = f"MCP工具执行失败: {str(e) or '未知错误，请检查服务器日志'}\n"
        yield json.dumps({
            "type": "message",
            "content": error_message,
            "timestamp": datetime.datetime.now().isoformat()
        })

